package com.zx.security.reactive.filter;

import com.zx.core.constant.TokenConstants;
import com.zx.core.context.TenantContextHolder;
import com.zx.security.reactive.service.ICustomTokenService;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.context.ReactiveSecurityContextHolder;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Mono;

import java.util.Objects;

/**
 * <p>
 * description: jwt_token 认证过滤器  <br>
 * create: 2024-08-16 00:59
 * </p>
 *
 * @author zhou  xun
 */

@Slf4j
@Component
public class JwtAuthenticationTokenFilter implements WebFilter {

    @Autowired
    ICustomTokenService iCustomTokenService;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
        String token = exchange.getRequest().getHeaders().getFirst(TokenConstants.AUTHENTICATION);
        String tokenWithoutPrefix = StringUtils.isNotEmpty(token) && token.startsWith(TokenConstants.PREFIX) ? token.substring(TokenConstants.PREFIX.length()) : token;
        // 如果token为空或无效，则直接继续处理
        if (StringUtils.isEmpty(tokenWithoutPrefix)) {
            return chain.filter(exchange);
        }

        // 验证token并获取用户信息
        return iCustomTokenService.getLoginUser(tokenWithoutPrefix)
                // 确保用户不为null
                .filter(Objects::nonNull)
                .flatMap(user -> {
                    UsernamePasswordAuthenticationToken auth = new UsernamePasswordAuthenticationToken(user, null, user.getAuthorities());

                    // 解析JWT token并提取租户信息
                    return iCustomTokenService.getReactiveJwtDecoder().decode(tokenWithoutPrefix)
                            .doOnNext(jwt -> {
                                auth.setDetails(jwt);

                                // 从JWT中提取租户信息并设置到上下文
                                String tenantId = jwt.getClaimAsString("tenantId");
                                String tenantName = jwt.getClaimAsString("tenantName");

                                if (StringUtils.isNotEmpty(tenantId)) {
                                    TenantContextHolder.setTenantId(tenantId);
                                }

                                if (StringUtils.isNotEmpty(tenantName)) {
                                    TenantContextHolder.setTenantName(tenantName);
                                }
                            })
                            .onErrorResume(throwable -> {
                                // 解析失败时忽略错误，继续处理
                                return Mono.empty();
                            })
                            .thenReturn(auth);
                })
                .defaultIfEmpty(null)
                .flatMap(auth -> {
                    if (auth != null) {
                        // 设置Authentication到SecurityContext
                        return chain.filter(exchange)
                                .contextWrite(ReactiveSecurityContextHolder.withAuthentication(auth));
                    } else {
                        return chain.filter(exchange);
                    }
                })
                .doFinally(signalType -> {
                    // 清除租户上下文
                    TenantContextHolder.clear();
                });
    }
}